import os
import json
import numpy as np
import pandas as pd
from collections import defaultdict
import time
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
import seaborn as sns
import pygame
from matplotlib.backends.backend_agg import FigureCanvasAgg
import io

# 文件路径
DATASET_PATH = 'assets/database/main_database.json'
OUTPUT_DIR = 'output'
FONT_PATH = 'assets/fonts/KeLaiYinLanDeYeWan-2.ttf'

# 加载字体
chinese_font = FontProperties(fname=FONT_PATH, size=24)
digit_font = FontProperties(fname=FONT_PATH, size=12)

def load_data(file_path):
    with open(file_path, 'r') as f:
        data = json.load(f)
    return data

def extract_features(data, username):
    current_user_data = defaultdict(list)
    all_users_data = defaultdict(list)
    user_keystrokes_count = defaultdict(int)
    
    # 用于临时存储每组击键数据的变量
    current_user_group_data = defaultdict(list)
    all_users_group_data = defaultdict(list)
    
    for entry in data:
        subject = entry['subject']
        keystrokes = entry['keystrokes']
        
        if user_keystrokes_count[subject] < 400:  # 只处理每个用户的前400次击键
            group_start_index = 0
            group_count = 0
            
            for i in range(len(keystrokes)):
                key = keystrokes[i]['key_content']
                press_time = keystrokes[i]['press_time']
                release_time = keystrokes[i]['release_time']
                duration = release_time - press_time
                
                # 记录持续时间
                all_users_data['durations'].append(duration)
                if subject == username:
                    current_user_data['durations'].append(duration)
                
                # 计算间隔时间
                if i > 0:
                    prev_release_time = keystrokes[i - 1]['release_time']
                    interval = press_time - prev_release_time
                    all_users_data['intervals'].append(interval)
                    if subject == username:
                        current_user_data['intervals'].append(interval)
                    
                    # 计算加速度(仅当interval不为0)
                    if interval != 0:
                        acceleration = (duration - (keystrokes[i - 1]['release_time'] - keystrokes[i - 1]['press_time'])) / interval
                        all_users_data['accelerations'].append(acceleration)
                        if subject == username:
                            current_user_data['accelerations'].append(acceleration)
                
                # 每10次击键计算一次平均速度
                if (i + 1) % 10 == 0 and i > 0:
                    group_press_time = keystrokes[group_start_index]['press_time']
                    group_release_time = keystrokes[i]['release_time']
                    group_duration = group_release_time - group_press_time
                    
                    if group_duration > 0:
                        avg_speed = 10 / group_duration  # 10次击键/总时间
                        
                        all_users_group_data['avg_speeds'].append(avg_speed)
                        if subject == username:
                            current_user_group_data['avg_speeds'].append(avg_speed)
                    
                    group_start_index = i + 1
                    group_count += 1
            
            user_keystrokes_count[subject] += len(keystrokes)
    
    # 将组平均速度合并到主数据中
    all_users_data['avg_speeds'] = all_users_group_data.get('avg_speeds', [])
    current_user_data['avg_speeds'] = current_user_group_data.get('avg_speeds', [])
    
    # 打印特征数量
    print(f"当前用户 {username} 的特征数量：")
    for feature, values in current_user_data.items():
        print(f"{feature}: {len(values)}")
    
    print("所有用户的特征数量：")
    for feature, values in all_users_data.items():
        print(f"{feature}: {len(values)}")
    
    return current_user_data, all_users_data

def create_output_folder(username):
    user_output_dir = os.path.join(OUTPUT_DIR, str(username))
    if not os.path.exists(user_output_dir):
        os.makedirs(user_output_dir)
    return user_output_dir

def plot_duration_histogram(current_user_data, output_dir):
    plt.figure(figsize=(12, 6))
    data = current_user_data['durations']
    n, bins, patches = plt.hist(data, bins=20, alpha=0.7, color='blue', edgecolor='black')
    
    for i in range(len(patches)):
        x = patches[i].get_x() + patches[i].get_width() / 2
        y = patches[i].get_height()
        if y > 0:
            plt.text(x, y, f'{int(y)}', ha='center', va='bottom', fontproperties=digit_font, color='black')

    for i in range(len(bins) - 1):
        bin_median = (bins[i] + bins[i + 1]) / 2
        plt.text(bin_median, -max(n) * 0.05, f'{bin_median:.3f}', ha='center', va='top', fontproperties=digit_font, color='black')

    plt.title('击键持续时间频率分布直方图', fontproperties=chinese_font)
    plt.xlabel('')
    plt.ylabel('频率', fontproperties=chinese_font)
    plt.xticks([])
    plt.tight_layout()
    
    # 保存图片
    plt.savefig(os.path.join(output_dir, '1_duration_histogram.png'))
    
    # 转换为Pygame可显示的Surface
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    img = pygame.image.load(buf)
    plt.close()
    
    return img

def plot_interval_histogram(current_user_data, output_dir):
    plt.figure(figsize=(12, 6))
    data = current_user_data['intervals']
    n, bins, patches = plt.hist(data, bins=20, alpha=0.7, color='green', edgecolor='black')
    
    for i in range(len(patches)):
        x = patches[i].get_x() + patches[i].get_width() / 2
        y = patches[i].get_height()
        if y > 0:
            plt.text(x, y, f'{int(y)}', ha='center', va='bottom', fontproperties=digit_font, color='black')

    for i in range(len(bins) - 1):
        bin_median = (bins[i] + bins[i + 1]) / 2
        plt.text(bin_median, -max(n) * 0.05, f'{bin_median:.3f}', ha='center', va='top', fontproperties=digit_font, color='black')

    plt.title('击键间隔时间频率分布直方图', fontproperties=chinese_font)
    plt.xlabel('')
    plt.ylabel('频率', fontproperties=chinese_font)
    plt.xticks([])
    plt.tight_layout()
    
    plt.savefig(os.path.join(output_dir, '2_interval_histogram.png'))
    
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    img = pygame.image.load(buf)
    plt.close()
    
    return img

def process_heatmap_data(current_user_data):
    return list(zip(current_user_data['durations'], current_user_data['intervals']))

def plot_heatmap(heatmap_data, output_dir):
    plt.figure(figsize=(10, 8))
    current_user_data = np.array(heatmap_data)
    sns.kdeplot(
        x=current_user_data[:, 0], 
        y=current_user_data[:, 1], 
        cmap="Reds", 
        fill=True,
    )

    plt.title('击键特征热力图', fontproperties=chinese_font)
    plt.xlabel('持续时间 (秒)', fontproperties=chinese_font)
    plt.ylabel('间隔时间 (秒)', fontproperties=chinese_font)
    plt.tight_layout()
    
    plt.savefig(os.path.join(output_dir, '3_heatmap.png'))
    
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    img = pygame.image.load(buf)
    plt.close()
    
    return img

def plot_radar_chart(current_user_data, all_users_data, output_dir, username):
    # 计算各特征的平均值
    all_users_avg = {
        '持续时间': np.mean(all_users_data['durations']),
        '间隔时间': np.mean(all_users_data['intervals']),
        '平均速度': np.mean(all_users_data['avg_speeds']),
        '加速度': None  # 先设为None，后面会计算
    }

    current_user_avg = {
        '持续时间': np.mean(current_user_data['durations']),
        '间隔时间': np.mean(current_user_data['intervals']),
        '平均速度': np.mean(current_user_data['avg_speeds']),
        '加速度': None  # 先设为None，后面会计算
    }

    # 改进加速度计算方法
    def calculate_improved_acceleration(keystrokes):
        if 'accelerations' not in keystrokes or len(keystrokes['accelerations']) == 0:
            return 0
            
        # 获取加速度数据
        accels = np.array(keystrokes['accelerations'])
        
        # 舍弃最大和最小的20%数据
        n = len(accels)
        keep_start = int(n * 0.2)
        keep_end = int(n * 0.8)
        
        if keep_end <= keep_start:
            return np.mean(accels)
            
        sorted_accels = np.sort(accels)
        trimmed_accels = sorted_accels[keep_start:keep_end]
        
        return np.mean(trimmed_accels)

    # 计算改进后的加速度
    all_users_avg['加速度'] = calculate_improved_acceleration(all_users_data)
    current_user_avg['加速度'] = calculate_improved_acceleration(current_user_data)

    # 获取所有特征的最大值和最小值用于标准化
    max_values = {
        '持续时间': max(all_users_avg['持续时间'], current_user_avg['持续时间']),
        '间隔时间': max(all_users_avg['间隔时间'], current_user_avg['间隔时间']),
        '平均速度': max(all_users_avg['平均速度'], current_user_avg['平均速度']),
        '加速度': max(abs(all_users_avg['加速度']), abs(current_user_avg['加速度']))
    }

    min_values = {
        '持续时间': min(all_users_avg['持续时间'], current_user_avg['持续时间']),
        '间隔时间': min(all_users_avg['间隔时间'], current_user_avg['间隔时间']),
        '平均速度': min(all_users_avg['平均速度'], current_user_avg['平均速度']),
        '加速度': -max_values['加速度']  # 使用对称范围
    }

    # 改进的标准化方法，确保所有值在0-1范围内
    def normalize(value, feature):
        # 对于加速度，我们使用对称范围
        if feature == '加速度':
            return (value - min_values[feature]) / (max_values[feature] - min_values[feature])
        else:
            return value / max_values[feature]

    # 准备标准化后的数据
    all_users_normalized = {
        '持续时间': normalize(all_users_avg['持续时间'], '持续时间'),
        '间隔时间': normalize(all_users_avg['间隔时间'], '间隔时间'),
        '平均速度': normalize(all_users_avg['平均速度'], '平均速度'),
        '加速度': normalize(all_users_avg['加速度'], '加速度')
    }

    current_user_normalized = {
        '持续时间': normalize(current_user_avg['持续时间'], '持续时间'),
        '间隔时间': normalize(current_user_avg['间隔时间'], '间隔时间'),
        '平均速度': normalize(current_user_avg['平均速度'], '平均速度'),
        '加速度': normalize(current_user_avg['加速度'], '加速度')
    }

    categories = list(all_users_normalized.keys())
    num_vars = len(categories)
    angles = np.linspace(0, 2 * np.pi, num_vars, endpoint=False).tolist()
    angles += angles[:1]

    fig, ax = plt.subplots(figsize=(8, 8), subplot_kw=dict(polar=True))
    
    # 设置雷达图的范围和网格
    ax.set_ylim(0, 1.0)  # 固定范围为0-1
    ax.set_rgrids([0.2, 0.4, 0.6, 0.8, 1.0], angle=0)

    # 绘制所有用户平均数据
    values = list(all_users_normalized.values())
    values += values[:1]
    ax.plot(angles, values, label='所有用户平均', color='blue', linewidth=2)
    ax.fill(angles, values, color='blue', alpha=0.1)

    # 绘制当前用户数据
    values = list(current_user_normalized.values())
    values += values[:1]
    ax.plot(angles, values, label=username, color='green', linewidth=2)  # 使用用户名作为图例
    ax.fill(angles, values, color='green', alpha=0.1)

    # 设置标签和标题
    ax.set_xticks(angles[:-1])
    ax.set_xticklabels(categories, fontproperties=chinese_font, size=12)
    ax.set_title('击键特征雷达图 (标准化)', fontproperties=chinese_font, size=14, y=1.1)
    
    # 添加图例
    ax.legend(loc='upper right', bbox_to_anchor=(1.2, 1.2), prop=chinese_font)

    # 添加刻度标签
    for label, angle in zip(ax.get_xticklabels(), angles[:-1]):
        if angle in (0, np.pi):
            label.set_horizontalalignment('center')
        elif 0 < angle < np.pi:
            label.set_horizontalalignment('left')
        else:
            label.set_horizontalalignment('right')

    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, '4_radar_chart.png'), dpi=120, bbox_inches='tight')
    
    buf = io.BytesIO()
    plt.savefig(buf, format='png', dpi=120, bbox_inches='tight')
    buf.seek(0)
    img = pygame.image.load(buf)
    plt.close()
    
    return img

def draw_colored_text(screen, font, text, position, colors, color_indices):
    x, y = position
    for i in range(len(text)):
        char = text[i]
        for start, end, color in color_indices:
            if start <= i < end:
                char_surface = font.render(char, True, color)
                screen.blit(char_surface, (x, y))
                x += char_surface.get_width()
                break
        else:
            char_surface = font.render(char, True, colors[0])
            screen.blit(char_surface, (x, y))
            x += char_surface.get_width()


def display_images(screen, images):
    """按顺序展示四张图片，支持导航功能"""
    pygame.display.set_caption("击键特征可视化")

    arrow_font = pygame.font.Font(FONT_PATH, 36)
    prompt_font = pygame.font.Font(FONT_PATH, 24)
    # 箭头符号
    left_arrow = arrow_font.render("←", True, (255, 0, 0))  # 左箭头为红色
    right_arrow = arrow_font.render("→", True, (0, 255, 0))  # 右箭头为绿色

    current_index = 0
    running = True

    while running:
        screen_width, screen_height = screen.get_size()

        # 清屏
        screen.fill((255, 255, 255))

        # 显示当前图片
        img = images[current_index]
        img_width, img_height = img.get_size()
        img_x = (screen_width - img_width) // 2
        img_y = (screen_height - img_height) // 2 - 30
        screen.blit(img, (img_x, img_y))

        # 显示导航提示
        prompt_y = screen_height - 50

        if current_index == 0:
            # 第一张图：只有下一张选项
            prompt_text = "按 右方向键 下一张图"
            colors = [(0, 0, 0), (0, 255, 0)]
            color_indices = [(2, 6, (0, 255, 0))]  # "右方向键"四个字变色(索引2-6)
            draw_colored_text(screen, prompt_font, prompt_text, (
                (screen_width - prompt_font.size(prompt_text)[0]) // 2, prompt_y), colors, color_indices)

            # 调整右箭头位置
            arrow_x = (screen_width - prompt_font.size(prompt_text)[0]) // 2 + prompt_font.size("按 ")[0]
            arrow_y = prompt_y - right_arrow.get_height() // 2
            screen.blit(right_arrow, (arrow_x, arrow_y))

        elif current_index == len(images) - 1:
            # 最后一张图：上一张和返回主菜单选项
            prompt_text = "按 左方向键 上一张图 按 回车键 返回主菜单"
            colors = [(0, 0, 0), (255, 0, 0), (0, 255, 0)]
            color_indices = [(2, 6, (255, 0, 0)), (13, 17, (0, 255, 0))]  # "左方向键"和"回车键"各四个字变色
            draw_colored_text(screen, prompt_font, prompt_text, (
                (screen_width - prompt_font.size(prompt_text)[0]) // 2, prompt_y), colors, color_indices)

            # 调整左箭头位置
            left_arrow_x = (screen_width - prompt_font.size(prompt_text)[0]) // 2 + prompt_font.size("按 ")[0]
            left_arrow_y = prompt_y - left_arrow.get_height() // 2
            screen.blit(left_arrow, (left_arrow_x, left_arrow_y))

        else:
            # 中间图片：上一张和下一张选项
            prompt_text = "按 左方向键 上一张图 按 右方向键 下一张图"
            colors = [(0, 0, 0), (255, 0, 0), (0, 255, 0)]
            color_indices = [(2, 6, (255, 0, 0)), (14, 18, (0, 255, 0))]  # "左方向键"和"右方向键"各四个字变色
            draw_colored_text(screen, prompt_font, prompt_text, (
                (screen_width - prompt_font.size(prompt_text)[0]) // 2, prompt_y), colors, color_indices)

            # 调整左箭头位置
            left_arrow_x = (screen_width - prompt_font.size(prompt_text)[0]) // 2 + prompt_font.size("按 ")[0]
            left_arrow_y = prompt_y - left_arrow.get_height() // 2
            screen.blit(left_arrow, (left_arrow_x, left_arrow_y))

            # 调整右箭头位置
            right_arrow_x = (screen_width - prompt_font.size(prompt_text)[0]) // 2 + prompt_font.size(
                "按 左方向键 上一张图 按 ")[0]
            right_arrow_y = prompt_y - right_arrow.get_height() // 2
            screen.blit(right_arrow, (right_arrow_x, right_arrow_y))

        pygame.display.flip()

        # 事件处理
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False

            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_RIGHT:
                    if current_index < len(images) - 1:
                        current_index += 1
                elif event.key == pygame.K_LEFT:
                    if current_index > 0:
                        current_index -= 1
                elif event.key == pygame.K_RETURN and current_index == len(images) - 1:
                    return True

        pygame.time.delay(10)

    return True
    

def visualize_keystroke_features(username, screen):
    """可视化击键特征主函数"""
    user_output_dir = create_output_folder(username)
    data = load_data(DATASET_PATH)
    current_user_data, all_users_data = extract_features(data, username)

    # 生成并保存图片，同时获取图片Surface对象
    images = []
    images.append(plot_duration_histogram(current_user_data, user_output_dir))
    images.append(plot_interval_histogram(current_user_data, user_output_dir))
    
    heatmap_data = process_heatmap_data(current_user_data)
    images.append(plot_heatmap(heatmap_data, user_output_dir))
    
    # 传入username参数
    images.append(plot_radar_chart(current_user_data, all_users_data, user_output_dir, username))

    # 在屏幕上显示图片
    return display_images(screen, images)

class KeystrokeVisualization:
    def __init__(self, screen):
        self.screen = screen
        self.font = pygame.font.Font(FONT_PATH, 36)
        self.input_name = ""
        self.input_name_active = True
        self.name_error_message = ""
        self.all_subjects = self.get_all_subjects()
        self.current_username = None
        self.visualization_complete = False
        self.should_return = False  

    def get_all_subjects(self):
        subjects = []
        if os.path.exists(DATASET_PATH):
            with open(DATASET_PATH, "r", encoding="utf-8") as file:
                data = json.load(file)
                seen = set()
                for user in data:
                    subject = str(user["subject"])
                    if subject not in seen:
                        subjects.append(subject)
                        seen.add(subject)
        return subjects

    def validate_name(self):
        if len(self.input_name) == 0:
            self.name_error_message = "用户名不能为空"
            return False
        elif self.input_name not in self.all_subjects:
            self.name_error_message = "用户不存在"
            return False
        return True

    def check_screen_size(self):
        """检测屏幕大小并调整字体大小"""
        screen_width, screen_height = self.screen.get_size()
        desktop_width, desktop_height = pygame.display.get_desktop_sizes()[0]
        
        if screen_width > desktop_width * 0.9 or screen_height > desktop_height * 0.9:
            self.font = pygame.font.Font(FONT_PATH, 54)
        else:
            self.font = pygame.font.Font(FONT_PATH, 36)

    def draw(self):
        self.check_screen_size()  # 每次绘制前检查屏幕大小
        self.screen.fill((255, 255, 255))
        screen_width, screen_height = self.screen.get_size()

        # 根据窗口大小动态调整字体
        is_fullscreen = screen_width > 1200  # 假设大于1200为全屏
        main_font_size = 36 if not is_fullscreen else 54
        main_font = pygame.font.Font(FONT_PATH, main_font_size)
        error_font_size = 24 if not is_fullscreen else 36
        error_font = pygame.font.Font(FONT_PATH, error_font_size)

        # 提示词位置
        prompt_text = main_font.render("请输入需要查询击键特征的用户名：", True, (0, 0, 0))
        prompt_rect = prompt_text.get_rect(center=(screen_width // 2, screen_height // 3))

        # 输入框位置
        input_width = int(screen_width * 0.6)
        input_height = int(screen_height * 0.08)
        input_rect = pygame.Rect(0, 0, input_width, input_height)
        input_rect.center = (screen_width // 2, screen_height // 2 - 30)

        # 按钮位置（下移并增加间距）
        button_width = int(screen_width * 0.15)
        button_height = int(screen_height * 0.08)
        confirm_rect = pygame.Rect(0, 0, button_width, button_height)
        confirm_rect.center = (screen_width // 2 - button_width - 20, screen_height // 2 + 100)
        
        back_rect = pygame.Rect(0, 0, button_width, button_height)
        back_rect.center = (screen_width // 2 + button_width + 20, screen_height // 2 + 100)

        # 错误提示位置（紧贴输入框下方）
        error_y = input_rect.bottom + 10

        # 绘制界面元素
        self.screen.blit(prompt_text, prompt_rect)
        pygame.draw.rect(self.screen, (200, 200, 200), input_rect)
        
        # 输入文本
        input_text = main_font.render(self.input_name, True, (0, 0, 0))
        self.screen.blit(input_text, (input_rect.x + 15, input_rect.y + (input_height - input_text.get_height()) // 2))

        # 按钮
        pygame.draw.rect(self.screen, (0, 200, 0), confirm_rect)
        confirm_text = main_font.render("确认", True, (0, 0, 0))
        self.screen.blit(confirm_text, confirm_text.get_rect(center=confirm_rect.center))

        pygame.draw.rect(self.screen, (200, 0, 0), back_rect)
        back_text = main_font.render("返回", True, (0, 0, 0))
        self.screen.blit(back_text, back_text.get_rect(center=back_rect.center))

        # 高亮效果
        mouse_pos = pygame.mouse.get_pos()
        if confirm_rect.collidepoint(mouse_pos):
            pygame.draw.rect(self.screen, (255, 255, 0), confirm_rect, 3)
        if back_rect.collidepoint(mouse_pos):
            pygame.draw.rect(self.screen, (255, 255, 0), back_rect, 3)

        # 错误提示
        if self.name_error_message:
            error_text = error_font.render(self.name_error_message, True, (255, 0, 0))
            error_rect = error_text.get_rect(midtop=(input_rect.centerx, error_y))
            self.screen.blit(error_text, error_rect)

        # 保存按钮rect
        self.confirm_button_rect = confirm_rect
        self.back_button_rect = back_rect

    def handle_event(self, event):
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_RETURN:
                if self.validate_name():
                    self.current_username = self.input_name
                    # 直接调用可视化函数并传入screen
                    visualize_keystroke_features(self.input_name, self.screen)
                    # 可视化完成后直接返回True，表示可以退出当前界面
                    self.should_return = True
            elif event.key == pygame.K_BACKSPACE:
                self.input_name = self.input_name[:-1]
                self.name_error_message = ""
            else:
                if event.unicode.isalnum():
                    self.input_name += event.unicode
                    self.name_error_message = ""
        
        # 新增：处理返回按钮点击
        if event.type == pygame.MOUSEBUTTONDOWN:
            if hasattr(self, 'back_button_rect') and self.back_button_rect.collidepoint(event.pos):
                self.should_return = True
            if hasattr(self, 'confirm_button_rect') and self.confirm_button_rect.collidepoint(event.pos):
                if self.validate_name():
                    self.current_username = self.input_name
                    visualize_keystroke_features(self.input_name, self.screen)
                    self.should_return = True
        return False

    def run(self):
        """主循环"""
        running = True
        while running and not self.should_return:  # 当需要返回时退出循环
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    return False
                self.handle_event(event)
            
            self.draw()
            pygame.display.flip()
        
        # 返回是否要退出当前界面
        return self.should_return